#include "problem.h"
#include "economy.h"
#include <strstream>
#include <limits.h>

const int kHoldoutNum = 1;

void InitProblem() {}

Problem::~Problem() {}

Thing Problem::ArgsType(Thing function_name, Thing args) {
  if (args.IsArray() && args.ArraySize() == 1)
    return ArgType(function_name, args.Ref(0));
  return NIL;
}

Thing Problem::ImplicitResult(Thing function_name, Thing arg) {
  Thing type = ArgType(function_name, arg);
  if (type == TRAINING_ARGUMENT) return Result(function_name, arg);
  if (type == TEST_ARGUMENT) return DefaultResult(function_name, arg);
  if (type == PREDICTION_ARGUMENT) return DefaultResult(function_name, arg);
  return NULL;
}

Thing NullToAbsentH(Thing x) {
  if (x.Is()) return x;
  return ABSENT;
}

Thing Problem::SystemCall
(EconomyThread *et, Syscall which, int32 num_parameters) {
  Thing ret = ExtensionSystemCall(et, which, num_parameters);
  if (ret != NOTIMPLEMENTED) return ret;
  Thing param0 = NIL, param1 = NIL, param2 = NIL;
  if (num_parameters > 0) param0 = et->vi_->GetParameter(0, num_parameters);
  if (num_parameters > 1) param1 = et->vi_->GetParameter(1, num_parameters);
  if (num_parameters > 2) param2 = et->vi_->GetParameter(2, num_parameters);
  // cout << "p0 = " << param0 << " p1 = " << param1 << " p2 = " << param2;
  switch (which) {
  case SYSCALL_PROBLEM_NUM_FUNCTIONS: 
    return Thing::MakeInteger(NumFunctions());
  case SYSCALL_PROBLEM_NTH_FUNCTION:
    if (!param0.IsInteger()) throw TYPE_EXCEPTION;
    return NthFunction(param0.GetInteger());
  case SYSCALL_PROBLEM_NUM_ARGS: 
    if (param1 == NIL) return NULL;
    return Thing::MakeInteger(NumArgs(param0, param1));
  case SYSCALL_PROBLEM_NTH_ARG:
    if (param2 == NIL) return NULL;
    if (!param1.IsInteger()) throw TYPE_EXCEPTION;
    return NthArg(param0, param1.GetInteger(), param2);
  case SYSCALL_PROBLEM_ARG_TYPE: 
    return ArgType(param0, param1);
  case SYSCALL_PROBLEM_IMPLICIT_RESULT: 
    return NullToAbsentH(ImplicitResult(param0, param1));
  case SYSCALL_PROBLEM_CORRECT_RESULT: 
    return NullToAbsentH(Result(param0, param1));
  case SYSCALL_PROBLEM_DESCRIPTION: 
    return Description(et, param0); 
  default: return NOTIMPLEMENTED;
  }
}

IntSequenceProblem::IntSequenceProblem() { 
  si_.Load(); 
  most_common_value_.clear();
  for (int64 i=0; i < si_.num_sequences_; i++) {
    hash_map<int64, int64, SimpleHash<int64> > counts;
    int64 most_common_value = 0;
    int64 max_count = 0;
    for (int64 j=0; j<NumTrainingArgs(i); j++) {
      int64 value = si_.sequences_[i][j];
      int64 count = ++(counts[value]);
      if (count > max_count) {
	max_count = count;
	most_common_value = value;
      }
    }
    most_common_value_.push_back(most_common_value);
  }
}

// Tests whether a function is in range, converts to
// an integer. Returns -1 on failure.
int64 IntSequenceProblem::GetID(Thing function_name) {
  if (!function_name.IsInteger()) return -1;
  int64 i = function_name.GetInteger();
  if (i<0 || i>= NumFunctions()) return -1;
  return i;
}

// assumes id is an integer in range. 
int64 IntSequenceProblem::NumTrainingArgs(int64 id) {
  return max(0, (si_.lengths_[id] - kHoldoutNum)); }

// The necessary methods
int64 IntSequenceProblem::NumFunctions() { 
  return si_.num_sequences_; }

Thing IntSequenceProblem::NthFunction(int64 n) {
  return Thing::MakeInteger(n); }

int64 IntSequenceProblem::NumArgs(Thing function_name, Thing arg_type) {
  int64 id = GetID(function_name); 
  if (id == -1) return 0;
  int64 all = si_.lengths_[id];
  if (arg_type == ALL_ARGUMENTS) return all;
  if (arg_type == TRAINING_ARGUMENT) return NumTrainingArgs(id);
  if (arg_type == TEST_ARGUMENT) return all - NumTrainingArgs(id);
  return 0;
}

// Implementation doesn't check out of range on argument
Thing IntSequenceProblem::NthArg(Thing function_name, int64 n, Thing arg_type) {
  /*
  cout << "::NthArg " 
       << "fname = " << function_name
       << " n = " << n
       << " arg_type = " << arg_type
       << endl;
  */
  if (arg_type != TEST_ARGUMENT) return Thing::MakeInteger(n);
  int64 id = GetID(function_name); 
  if (id == -1) return NULL;
  return Thing::MakeInteger(n + NumTrainingArgs(id)); 
}

// Does range checks
Thing IntSequenceProblem::ArgType(Thing function_name, Thing arg) {
  int64 id = GetID(function_name); 
  if (id == -1) return NIL;
  if (!(arg.IsInteger())) return NIL;
  int64 i_arg = arg.GetInteger();
  if (i_arg < 0 || i_arg >= si_.lengths_[id]) return NIL;
  if (i_arg < NumTrainingArgs(id)) return TRAINING_ARGUMENT;
  return TEST_ARGUMENT;
}

Thing IntSequenceProblem::Description
(EconomyThread * et, Thing function_name) {
  if (!(function_name.IsInteger())) return NIL;
  int64 n = function_name.GetInteger();
  if (n < 0 || n >= NumFunctions()) return NIL;
  char * desc = si_.GetDescription(n);
  Thing ret = CStringToThingString(et->a_, desc);
  return ret;
}

Thing IntSequenceProblem::Result(Thing function_name, Thing arg) {
  Thing arg_type = ArgType(function_name, arg);
  if (arg_type == NIL) return NULL;
  int64 id = GetID(function_name);
  int64 i_arg = arg.GetInteger();
  return Thing::MakeInteger(si_.sequences_[id][i_arg]);
}

Thing IntSequenceProblem::DefaultResult(Thing function_name, Thing arg) {
  return Thing::MakeInteger(most_common_value_[function_name.GetInteger()])
;
  // return T;
  // return Thing::MakeInteger(-1);
}


void ExplicitProblem::AddTuple(Thing tuple) {
  // cout << "Adding triple " << tuple << endl;
  CHECK(tuple.IsArray());
  CHECK(tuple.ArraySize() >= 3); // typically 5, but 3 for function description
  Thing fname = tuple.Ref(0);
  Thing arg = tuple.Ref(1);
  Thing arg_type = tuple.Ref(2);
  if (not (fnames_set_ % fname.ID())) {
    fnames_set_.insert(fname.ID());
    fnames_.push_back(fname);
  }
  if (arg == PROBLEM_DESCRIPTION) {
    fname_to_description_[fname.ID()] = arg_type;
    return;
  }
  CHECK(tuple.ArraySize() == 5);
  fname_type_to_args_[Thing::MakeArrayID2(fname, arg_type)].push_back(arg);
  fname_type_to_args_[Thing::MakeArrayID2(fname, ALL_ARGUMENTS)].push_back(arg);
  uint64 fname_arg = Thing::MakeArrayID2(fname, arg);
  if (fname_arg_to_tuple_ % fname_arg) {
    cout << "duplicate fname_arg " << tuple << endl;
    cout << "previous result " 
	 << fname_arg_to_tuple_[fname_arg] << endl;
    CHECK(false);
  }
  fname_arg_to_tuple_[fname_arg] = tuple;
}


void ExplicitProblem::LoadFromFile(string fn) {
  ifstream input(fn.c_str());
  ThreadAllocator a;
  Vinterpreter vi;
  vi.a_ = &a;  
  Thing code;
  int64 counter = 0;
  while (1) {
    counter++;
    if (counter % 10000 == 0) cout << "Loaded " << counter << " tuples\n";   
    code = ReadThing(&a, input);
    if (code.IsNot()) return;
    Thing result = vi.SafeExecute(code, (int64) 1e9, kNoSystemMode, NIL);
    result = a.Globalize(result);
    if (result.Is() && result.IsArray()) {
      for (uint64 i=0; i<result.ArraySize(); i++)
	AddTuple(result.Ref(i));
    }
  }
}

ExplicitProblem::ExplicitProblem(string fns) {
  for (uint i=0; i<fns.size(); i++) if (fns[i] == ',') fns[i] = ' ';
  istringstream istr(fns);
  string fn;
  while (istr >> fn) LoadFromFile(fn);
}

int64 ExplicitProblem::NumFunctions() {
  return fnames_.size(); }

Thing ExplicitProblem::NthFunction(int64 n) {
  if (n < 0 || n >= NumFunctions()) return NULL;
  return fnames_[n];
}

int64 ExplicitProblem::NumArgs(Thing fname, Thing arg_type) {
  uint64 fname_type = Thing::MakeArrayID2(fname, arg_type);
  vector<Thing> * args = fname_type_to_args_  % fname_type;
  if (!args) return 0;
  return args->size();
}

Thing ExplicitProblem::NthArg(Thing fname, int64 n, Thing arg_type) {
  uint64 fname_type = Thing::MakeArrayID2(fname, arg_type);
  vector<Thing> * args = fname_type_to_args_ % fname_type;
  if (!args) return NULL;
  if (n < 0 || (n >= (int64)args->size())) return NULL;
  return (*args)[n];
}

Thing ExplicitProblem::Description
(EconomyThread * et, Thing function_name) {
  Thing * look = fname_to_description_ % function_name.ID();
  if (!look) return NIL;
  return *look;
}

Thing ExplicitProblem::ArgType(Thing fname, Thing arg){
  Thing * look = fname_arg_to_tuple_ % Thing::MakeArrayID2(fname, arg);
  if (!look) return NIL;
  // cout << "Look = " << *look << endl;
  CHECK(look->IsArray());
  return look->Ref(2);
}
Thing ExplicitProblem::Result(Thing fname, Thing arg){
  Thing * look = fname_arg_to_tuple_ % Thing::MakeArrayID2(fname, arg);
  if (!look) return NULL;
  return look->Ref(3);
}
Thing ExplicitProblem::DefaultResult(Thing fname, Thing arg){
  Thing * look = fname_arg_to_tuple_ % Thing::MakeArrayID2(fname, arg);
  if (!look) return NULL;
  return look->Ref(4);
}


SequenceProblem::SequenceProblem() {
  ds_ = new DataSet(); ds_->Load(); }

uint64 SequenceProblem::TokenCount(uint64 token_id) {
  CHECK(token_id < Lex()->Size());
  uint64 *begin = (uint64*) ds_->token_counts_.begin();
  return begin[token_id];
}

// in micronats
int64 SequenceProblem::TokenLikelihood(uint64 token_id) {
  CHECK(token_id < Lex()->Size());
  uint64 *begin = (uint64*) ds_->token_counts_.begin();
  return 1 * (MLn(begin[token_id]) - MLn(GetCorpus()->size()));
}

Thing SequenceProblem::NthFunction(int64 n) {
  CHECK(n==0); return CORPUS; }

int64 SequenceProblem::NumArgs(Thing function_name, Thing arg_type) {
  int64 ret = GetCorpus()->size();
  if (function_name != CORPUS) return 0;
  if ((arg_type == ALL_ARGUMENTS) || (arg_type == TRAINING_ARGUMENT))
    return ret;
  return 0; 
}

Thing SequenceProblem::NthArg(Thing function_name, int64 n, Thing arg_type) {
  if (function_name != CORPUS) return NULL;
  if ((arg_type != ALL_ARGUMENTS) && (arg_type != TRAINING_ARGUMENT))
    return NULL;
  if (n < 0) return NULL;
  if (n >= (int64)GetCorpus()->size()) return NULL;
  return Thing::MakeInteger(n);
}

Thing SequenceProblem::ArgType(Thing function_name, Thing arg) {
  if (arg.GetType() != INTEGER) return NIL;
  int64 arg_i = arg.GetInteger();
  Thing nth_arg = NthArg(function_name, arg_i, ALL_ARGUMENTS);
  if (nth_arg.IsNot()) return NIL;
  return TRAINING_ARGUMENT;
}

Thing SequenceProblem::Result(Thing function_name, Thing arg) {
  if (function_name != CORPUS) return NULL;
  if (!arg.IsInteger()) return NULL;
  int64 i_pos = arg.GetInteger();
  if (i_pos < 0 || i_pos >= (int64)GetCorpus()->size()) return NULL;
  return Thing::MakeToken((*GetCorpus())[i_pos]);
}

Thing SequenceProblem::ExtensionSystemCall
(EconomyThread *et, Syscall which, int32 num_parameters) {
  Thing param0 = NIL;
  if (num_parameters > 0) param0 = et->vi_->GetParameter(0, num_parameters);
  switch (which) {
  // case SYSCALL_TOKENIZE: return TokenizeCall(et, param0);
  case SYSCALL_SEARCH: return SearchCall(et, param0);
  case SYSCALL_SUFFIX: return SuffixCall(param0);
  // case SYSCALL_LEXICONSIZE: return LexiconSizeCall(param0);
  default: return NOTIMPLEMENTED;
  }
}

// (search '(tok0 tok1 ... )) -> (suffix-start suffix-length)
// TODO: use CPU here
Thing SequenceProblem::SearchCall (EconomyThread * et, Thing param) { 
  // cout << "::SearchCall " << param.ToString() << endl;
  if (!param.IsArray()) throw TYPE_EXCEPTION;
  uint32 uint32_buffer[1000];
  Box *b = param.GetBox();
  uint64 size = b->ArraySize();
  for (uint64 i=0; i<size; i++) {
    Thing t = b->Ref(i);
    if (!t.IsToken()) throw TYPE_EXCEPTION;
    uint32_buffer[i] = t.GetUnboxedValue();
  }
  pair<uint64, uint64> look =
    DS()->index_.Search(uint32_buffer, size);
  Thing s_start = Thing::MakeInteger(look.first);
  Thing s_length = Thing::MakeInteger(look.second);
  return et->a_->MakeImmutableArray2(LOCAL_ALLOCATION, s_start, s_length);
}

// (suffix <suffix position>)
//    the position in the data that is the nth suffix
//    in the suffix sort
Thing SequenceProblem::SuffixCall(Thing param) {
  // if (et_->in_test_) return NIL;
  if (!param.IsInteger()) return NIL;
  int64 posi = param.GetInteger();
  if (posi < 0) return NIL;
  uint64 pos = (uint64) posi;
  SuffixIndex &si = DS()->index_;
  uint64 size = si.index_.size();
  if (pos >= size) return NIL;
  int64 ret = si.index_[pos];
  return Thing::MakeInteger(ret);
}

/*

int64 MetaProblem::NumFunctions() {
  int64 ret = 0;
  for (uint64 i=0; i<problems_.size(); i++) 
    ret += problems_[i]->NumFunctions();
  return ret; }

Thing MetaProblem::NthFunction(int64 n) {
  if (n < 0 || n >= NumFunctions()) throw OUT_OF_RANGE_EXCEPTION;
  for (uint64 i=0; i<problems_.size(); i++) {
    int64 nf = problems_[i]->NumFunctions();
    if (n < nf) return problems_[i]->NthFunction(n);
    n -= nf;
  }
  CHECK(false);
  return NULL;
}

int64 MetaProblem::NumArgs(Thing function_name) {
  int64 ret = 0;
  for (uint64 i=0; i<problems_.size(); i++)     
    ret += problems_[i]->NumArgs(function_name);
  return ret;
}

int64 MetaProblem::NumTestArgs(Thing function_name) {
  int64 ret = 0;
  for (uint64 i=0; i<problems_.size(); i++)     
    ret += problems_[i]->NumTestArgs(function_name);
  return ret;
}


Thing MetaProblem::NthArg(Thing function_name, int64 n) {
  if (n < 0) throw OUT_OF_RANGE_EXCEPTION;
  for (uint64 i=0; i<problems_.size(); i++) {
    int64 na = problems_[i]->NumArgs(function_name);
    if (n < na) return problems_[i]->NthArg(function_name, n);
    n -= na;
  }
  throw OUT_OF_RANGE_EXCEPTION;
}

Thing MetaProblem::NthTestArg(Thing function_name, int64 n) {
  if (n < 0) throw OUT_OF_RANGE_EXCEPTION;
  for (uint64 i=0; i<problems_.size(); i++) {
    int64 na = problems_[i]->NumTestArgs(function_name);
    if (n < na) return problems_[i]->NthTestArg(function_name, n);
    n -= na;
  }
  throw OUT_OF_RANGE_EXCEPTION;
}

Thing MetaProblem::Result(Thing function_name, Thing arg) {
  for (uint64 i=0; i<problems_.size(); i++) {
    Thing ret = problems_[i]->Result(function_name, arg);
    if (ret.Is()) return ret;
  }
  return NULL;
}

Thing MetaProblem::TestResult(Thing function_name, Thing arg) {
  for (uint64 i=0; i<problems_.size(); i++) {
    Thing ret = problems_[i]->TestResult(function_name, arg);
    if (ret.Is()) return ret;
  }
  return NULL;
}

void MetaProblem::AddProblem(Problem *p) {
  problems_.push_back(p);
}

Thing MetaProblem::SpecificProblemSystemCall(EconomyThread *et, 
					     Syscall which, 
					     int32 num_parameters) {
  for (uint64 i=0; i<problems_.size(); i++) {
    Thing val 
      = problems_[i]->SpecificProblemSystemCall(et, which, num_parameters);
    if (val != NOTIMPLEMENTED) return val;
  }
  return NOTIMPLEMENTED;
}

Thing MetaProblem::ProblemDescription
(EconomyThread *et, Thing function_name) {
  for (uint64 i=0; i<problems_.size(); i++) {
    Thing val 
      = problems_[i]->ProblemDescription(et, function_name);
    if (val != NIL) return val;
  }
  return NIL;
}

Thing SequenceProblem::SpecificProblemSystemCall
(EconomyThread *et, Syscall which, int32 num_parameters) {
  Thing param0 = NIL;
  if (num_parameters > 0) param0 = et->vi_->GetParameter(0, num_parameters);
  switch (which) {
  case SYSCALL_TOKENIZE: return TokenizeCall(et, param0);
  case SYSCALL_SEARCH: return SearchCall(et, param0);
  case SYSCALL_SUFFIX: return SuffixCall(param0);
  case SYSCALL_LEXICONSIZE: return LexiconSizeCall(param0);
  default: return NOTIMPLEMENTED;
  }
}

Thing SequenceProblem::TokenizeCall(EconomyThread * et, Thing param) {
  if (!param.IsArray()) return NIL;
  vector<char> v;
  ArrayThingToVectorChars(param, &v);
  v.push_back('\0');
  XMLTokenizer t;
  t.Init(&v[0], &v[0] + v.size() - 1);
  vector<Thing > output;
  Token *tok;
  Lexicon * lex = Lex();
  while ((tok = t.GetNextToken())) {
    uint32 n = lex->StringToID(tok->start_, tok->end_);
    if (n == lex->Size()) return NIL;
    output.push_back(Thing::MakeToken(n));
  }
  Thing ret = et->a_->MakeImmutableArray(LOCAL_ALLOCATION, output);
  return ret;
}

// (search '(tok0 tok1 ... )) -> (suffix-start suffix-length)
// TODO: use CPU here
Thing SequenceProblem::SearchCall (EconomyThread * et, Thing param) { 
  // cout << "::SearchCall " << param.ToString() << endl;
  if (!param.IsArray()) throw TYPE_EXCEPTION;
  uint32 uint32_buffer[1000];
  Box *b = param.GetBox();
  uint64 size = b->ArraySize();
  for (uint64 i=0; i<size; i++) {
    Thing t = b->Ref(i);
    if (!t.IsToken()) throw TYPE_EXCEPTION;
    uint32_buffer[i] = t.GetUnboxedValue();
  }
  pair<uint64, uint64> look =
    DS()->index_.Search(uint32_buffer, size);
  Thing s_start = Thing::MakeInteger(look.first);
  Thing s_length = Thing::MakeInteger(look.second);
  return et->a_->MakeImmutableArray2(LOCAL_ALLOCATION, s_start, s_length);
}

// (suffix <suffix position>)
//    the position in the data that is the nth suffix
//    in the suffix sort
Thing SequenceProblem::SuffixCall(Thing param) {
  // if (et_->in_test_) return NIL;
  if (!param.IsInteger()) return NIL;
  int64 posi = param.GetInteger();
  if (posi < 0) return NIL;
  uint64 pos = (uint64) posi;
  SuffixIndex &si = DS()->index_;
  uint64 size = si.index_.size();
  if (pos >= size) return NIL;
  int64 ret = si.index_[pos];
  return Thing::MakeInteger(ret);
}

Thing SequenceProblem::LexiconSizeCall(Thing param) {
  return Thing::MakeInteger(Lex()->Size()); }


*/
